#!/bin/bash
# AWS停止脚本 - 停止AWS节点上的Turritopsis协议

set -e

# 配置参数
KEY_NAME="turritopsis-key"
PROJECT_DIR="/home/ec2-user/turritopsis"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查hosts文件
check_hosts_file() {
    if [ ! -f "aws_hosts.txt" ]; then
        log_error "aws_hosts.txt not found. Please run aws_deploy.sh first."
        exit 1
    fi
    
    log_info "Found $(wc -l < aws_hosts.txt) nodes in aws_hosts.txt"
}

# 停止所有进程
stop_all_processes() {
    log_info "Stopping all Turritopsis processes on all nodes..."
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Stopping processes on $node_name ($public_ip)..."
            
            # 创建停止脚本
            cat > "stop_node_${node_id}.sh" << 'EOF'
#!/bin/bash
cd $PROJECT_DIR

# 停止Turritopsis进程
if [ -f turritopsis.pid ]; then
    PID=$(cat turritopsis.pid)
    if ps -p $PID > /dev/null 2>&1; then
        kill $PID
        echo "Stopped Turritopsis (PID: $PID)"
    else
        echo "Turritopsis process not running"
    fi
    rm -f turritopsis.pid
else
    echo "No Turritopsis PID file found"
fi

# 停止路由器进程
if [ -f router.pid ]; then
    PID=$(cat router.pid)
    if ps -p $PID > /dev/null 2>&1; then
        kill $PID
        echo "Stopped TCP router (PID: $PID)"
    else
        echo "Router process not running"
    fi
    rm -f router.pid
else
    echo "No router PID file found"
fi

# 停止ACS服务器进程
if [ -f acs_server.pid ]; then
    PID=$(cat acs_server.pid)
    if ps -p $PID > /dev/null 2>&1; then
        kill $PID
        echo "Stopped ACS server (PID: $PID)"
    else
        echo "ACS server process not running"
    fi
    rm -f acs_server.pid
else
    echo "No ACS server PID file found"
fi

# 强制杀死任何残留的turritopsis进程
pkill -f turritopsis_full 2>/dev/null || true
pkill -f acs_stub_server 2>/dev/null || true

echo "All processes stopped on $(hostname)"
EOF
            
            # 复制并执行停止脚本
            scp -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no "stop_node_${node_id}.sh" ec2-user@$public_ip:/tmp/
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "chmod +x /tmp/stop_node_${node_id}.sh && PROJECT_DIR=$PROJECT_DIR /tmp/stop_node_${node_id}.sh"
            
            ((node_id++))
        fi
    done < aws_hosts.txt
    
    log_info "All processes stopped on all nodes"
}

# 检查进程状态
check_process_status() {
    log_info "Checking process status on all nodes..."
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Checking $node_name ($public_ip)..."
            
            # 检查进程状态
            status=$(ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no -o ConnectTimeout=5 ec2-user@$public_ip \
                "cd $PROJECT_DIR && \
                turritopsis_status=\$(if [ -f turritopsis.pid ]; then ps -p \$(cat turritopsis.pid) > /dev/null 2>&1 && echo 'RUNNING' || echo 'STOPPED'; else echo 'NO_PID'; fi); \
                router_status=\$(if [ -f router.pid ]; then ps -p \$(cat router.pid) > /dev/null 2>&1 && echo 'RUNNING' || echo 'STOPPED'; else echo 'NO_PID'; fi); \
                acs_status=\$(if [ -f acs_server.pid ]; then ps -p \$(cat acs_server.pid) > /dev/null 2>&1 && echo 'RUNNING' || echo 'STOPPED'; else echo 'NO_PID'; fi); \
                echo \"Turritopsis: \$turritopsis_status, Router: \$router_status, ACS: \$acs_status\"" 2>/dev/null || echo "ERROR")
            
            echo "Node $node_id ($public_ip): $status"
            ((node_id++))
        fi
    done < aws_hosts.txt
}

# 清理日志文件
cleanup_logs() {
    log_info "Cleaning up log files on all nodes..."
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Cleaning logs on $node_name ($public_ip)..."
            
            # 清理日志文件
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "cd $PROJECT_DIR && \
                rm -f *.log *.pid && \
                echo 'Logs cleaned on $(hostname)'"
            
            ((node_id++))
        fi
    done < aws_hosts.txt
    
    log_info "Log cleanup completed on all nodes"
}

# 收集最终指标
collect_final_metrics() {
    log_info "Collecting final metrics before cleanup..."
    
    # 运行指标收集脚本
    if [ -f "aws_collect_metrics.sh" ]; then
        ./aws_collect_metrics.sh collect
        log_info "Final metrics collected"
    else
        log_warn "aws_collect_metrics.sh not found, skipping final metrics collection"
    fi
}

# 清理临时文件
cleanup_temp_files() {
    log_info "Cleaning up temporary files..."
    
    # 删除本地临时脚本
    rm -f stop_node_*.sh
    
    log_info "Temporary files cleaned up"
}

# 主函数
main() {
    log_info "Stopping Turritopsis on AWS nodes..."
    
    check_hosts_file
    collect_final_metrics
    stop_all_processes
    check_process_status
    cleanup_logs
    cleanup_temp_files
    
    log_info "Turritopsis stopped on all nodes!"
    log_info "Final metrics have been collected"
}

# 脚本参数处理
case "${1:-}" in
    "stop")
        main
        ;;
    "status")
        check_hosts_file
        check_process_status
        ;;
    "cleanup")
        check_hosts_file
        cleanup_logs
        cleanup_temp_files
        log_info "Cleanup completed"
        ;;
    "force-stop")
        log_warn "Force stopping all processes..."
        check_hosts_file
        
        node_id=0
        while read -r node_name public_ip private_ip; do
            if [[ "$node_name" =~ ^node- ]]; then
                log_info "Force stopping $node_name ($public_ip)..."
                ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                    "pkill -9 -f turritopsis_full; pkill -9 -f acs_stub_server; echo 'Force stopped on $(hostname)'"
                ((node_id++))
            fi
        done < aws_hosts.txt
        
        log_info "Force stop completed"
        ;;
    *)
        echo "Usage: $0 {stop|status|cleanup|force-stop}"
        echo "  stop       - Stop all processes gracefully"
        echo "  status     - Check process status"
        echo "  cleanup    - Clean up log files"
        echo "  force-stop - Force kill all processes"
        exit 1
        ;;
esac